from fastapi import FastAPI, Depends, HTTPException, UploadFile, File, Form, Header, Request
from sqlalchemy.orm import Session
from database import SessionLocal, engine
from models import Base, User, History, DetectionRecord, Task, DetectionDetail, VehicleDetail, TaskStatistics, TASK_STATUS_PENDING, TASK_STATUS_PROCESSING, TASK_STATUS_COMPLETED, TASK_STATUS_FAILED
import shutil
import os
from datetime import datetime, timedelta
from ultralytics import YOLO
import cv2
import json
from typing import List, Optional, Dict
from celery_config import celery_app
from celery.result import AsyncResult
import uuid
from sqlalchemy import Column, String, Integer, Float, DateTime, ForeignKey
from sqlalchemy.orm import relationship
from pydantic import BaseModel
from fastapi.responses import FileResponse
import logging
import base64
from fastapi.staticfiles import StaticFiles
import socket
from contextlib import closing

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 获取当前文件所在目录的绝对路径
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
# 构建上传文件目录路径
UPLOAD_DIR = os.path.join(os.path.dirname(BASE_DIR), "uploads")
# 确保上传目录存在
os.makedirs(UPLOAD_DIR, exist_ok=True)

def normalize_path(path):
    """将Windows路径转换为URL格式的路径"""
    if path:
        # 将反斜杠转换为正斜杠
        path = path.replace('\\', '/')
        # 确保路径以/uploads开头
        if not path.startswith('/uploads/'):
            # 获取相对于UPLOAD_DIR的路径
            rel_path = os.path.relpath(path, UPLOAD_DIR).replace('\\', '/')
            path = f'/uploads/{rel_path}'
    return path

# 导入任务
from tasks import process_task

Base.metadata.create_all(bind=engine)
app = FastAPI()

# 加载模型（只需加载一次）
MODEL_PATH = os.path.join(os.path.dirname(BASE_DIR), "Models-----YOLO11", "yolo11n.pt")
yolo_model = YOLO(MODEL_PATH)

def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

# 用户注册
@app.post("/register")
def register(username: str = Form(...), password: str = Form(...)):
    db = SessionLocal()
    try:
        # 检查用户名是否已存在
        if db.query(User).filter(User.username == username).first():
            raise HTTPException(status_code=400, detail="用户名已存在")
        
        # 创建新用户
        user = User(username=username, password=password)
        db.add(user)
        db.commit()
        db.refresh(user)
        
        return {"message": "注册成功", "user_id": user.id}
    except HTTPException as e:
        raise e
    except Exception as e:
        db.rollback()
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        db.close()

# 添加获取当前用户的函数
async def get_current_user(
    authorization: Optional[str] = Header(None),
    db: Session = Depends(get_db)
) -> User:
    if not authorization:
        raise HTTPException(
            status_code=401,
            detail="未提供认证信息"
        )
    
    try:
        # 从请求头中获取用户ID
        user_id = int(authorization)
        user = db.query(User).filter(User.id == user_id).first()
        if not user:
            raise HTTPException(
                status_code=401,
                detail="用户不存在"
            )
        return user
    except ValueError:
        raise HTTPException(
            status_code=401,
            detail="无效的认证信息"
        )

# 修改登录接口，返回token
@app.post("/login")
def login(username: str = Form(...), password: str = Form(...), db: Session = Depends(get_db)):
    user = db.query(User).filter(User.username == username, User.password == password).first()
    if not user:
        raise HTTPException(status_code=400, detail="用户名或密码错误")
    return {"success": True, "user_id": user.id, "token": str(user.id)}  # 使用用户ID作为token

# 修改密码
@app.put("/users/{user_id}/password")
def change_password(user_id: int, new_password: str = Form(...), db: Session = Depends(get_db)):
    user = db.get(User, user_id)
    if not user:
        raise HTTPException(status_code=404, detail="用户不存在")
    user.password = new_password
    db.commit()
    return {"success": True}

# 注销
@app.delete("/users/{user_id}")
def delete_user(user_id: int, db: Session = Depends(get_db)):
    user = db.get(User, user_id)
    if not user:
        raise HTTPException(status_code=404, detail="用户不存在")
    # 级联删除该用户的所有任务及相关数据
    # 删除任务相关的详情、车辆详情、统计
    user_tasks = db.query(Task).filter(Task.user_id == user_id).all()
    for task in user_tasks:
        db.query(DetectionDetail).filter(DetectionDetail.task_id == task.id).delete()
        db.query(VehicleDetail).filter(VehicleDetail.task_id == task.id).delete()
        db.query(TaskStatistics).filter(TaskStatistics.task_id == task.id).delete()
    # 删除任务
    db.query(Task).filter(Task.user_id == user_id).delete()
    # 删除历史记录
    db.query(History).filter(History.user_id == user_id).delete()
    # 删除检测记录（如果有 DetectionRecord 表）
    db.query(DetectionRecord).filter(DetectionRecord.user_id == user_id).delete()
    # 最后删除用户
    db.delete(user)
    db.commit()
    return {"success": True}

def is_image(filename):
    return filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))

def is_video(filename):
    return filename.lower().endswith(('.mp4', '.avi', '.mov', '.mkv', '.wmv'))

# 任务管理
@app.post("/upload")
async def upload_file(
    files: List[UploadFile] = File(...),
    user_id: int = Form(...),
    conf: float = Form(0.25),      # 置信度阈值
    iou: float = Form(0.45),       # IoU阈值
    max_det: int = Form(300),      # 最大检测数
    db: Session = Depends(get_db)
):
    logger.info(f"\n{'='*50}")
    logger.info(f"[{datetime.now()}] Starting multiple files upload")
    logger.info(f"User ID: {user_id}")
    logger.info(f"Number of files: {len(files)}")
    logger.info(f"Parameters: conf={conf}, iou={iou}, max_det={max_det}")
    logger.info(f"Upload directory: {UPLOAD_DIR}")
    
    task_infos = []
    
    for file in files:
        logger.info(f"\nProcessing file: {file.filename}")
        
        # 生成任务ID
        task_id = str(uuid.uuid4())
        logger.info(f"Generated task ID: {task_id}")
        
        # 创建用户目录
        user_dir = os.path.join(UPLOAD_DIR, f"user_{user_id}")
        os.makedirs(user_dir, exist_ok=True)
        
        # 使用绝对路径保存文件
        timestamp = datetime.now().strftime('%Y%m%d%H%M%S')
        input_filename = f"input_{timestamp}_{file.filename}"
        input_path = os.path.join(user_dir, input_filename)
        
        # 保存文件
        try:
            contents = await file.read()
            with open(input_path, "wb") as f:
                f.write(contents)
            logger.info(f"File saved to: {input_path}")
        except Exception as e:
            logger.error(f"Error saving file: {str(e)}")
            raise HTTPException(status_code=500, detail=f"Failed to save uploaded file {file.filename}: {str(e)}")
        finally:
            await file.close()
        
        # 验证文件是否成功保存
        if not os.path.exists(input_path):
            logger.error(f"Error: Failed to save file to {input_path}")
            raise HTTPException(status_code=500, detail=f"Failed to save uploaded file {file.filename}")

        # 确定媒体类型
        media_type = "image" if file.filename.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.gif')) else "video"

        # 创建任务记录
        task = Task(
            id=task_id,
            user_id=user_id,
            input_file=input_path,
            status=TASK_STATUS_PENDING,
            media_type=media_type,
            confidence_threshold=conf,
            iou_threshold=iou,
            max_detections=max_det,
            pedestrian_count=0,
            vehicle_count=0
        )
        db.add(task)
        
        db.commit()
        db.refresh(task)

        # Dispatch the task to Celery
        logger.info(f"Dispatching task {task_id} for file {file.filename}")
        try:
            # 调用 Celery 任务
            process_task.delay(task_id, input_path, conf, iou, max_det)
            logger.info(f"Successfully dispatched task {task_id} to Celery")
        except Exception as e:
            logger.error(f"Failed to dispatch task to Celery: {str(e)}")
            task.status = TASK_STATUS_FAILED
            task.error = f"Failed to dispatch task to Celery: {str(e)}"
            db.commit()
            raise HTTPException(status_code=500, detail=f"Failed to process file {file.filename}: {str(e)}")
        
        # 构建任务信息
        task_info = {
            "task_id": task_id,
            "filename": file.filename,
            "media_type": media_type
        }
        task_infos.append(task_info)
    
    logger.info(f"{'='*50}\n")
    return {
        "success": True,
        "message": f"Successfully uploaded {len(files)} files",
        "tasks": task_infos
    }

# 定义请求模型
class BatchOperationRequest(BaseModel):
    operation_type: str
    user_id: int
    selection_type: str = "selected"
    history_ids: Optional[List[int]] = None

# 定义结果查询请求模型
class ResultQueryRequest(BaseModel):
    task_ids: List[str]

# 批量获取处理结果
@app.post("/result")
async def result(request: ResultQueryRequest, db: Session = Depends(get_db)):
    try:
        results = []
        errors = []
        total_tasks = len(request.task_ids)
        successful_tasks = 0
        failed_tasks = 0
        
        for task_id in request.task_ids:
            try:
                task = db.query(Task).filter(Task.id == task_id).first()
                if not task:
                    errors.append(f"Task {task_id} not found")
                    failed_tasks += 1
                    continue
                
                # 获取用户信息
                user = db.query(User).filter(User.id == task.user_id).first()
                if not user:
                    errors.append(f"User not found for task {task_id}")
                    failed_tasks += 1
                    continue
                
                # 获取检测记录
                detection_record = db.query(DetectionRecord).filter(
                    DetectionRecord.user_id == task.user_id,
                    DetectionRecord.image_path == task.input_file
                ).first()
                
                # 获取任务统计信息
                statistics = None
                if task.statistics:
                    stats = task.statistics[0]
                    statistics = {
                        "max_confidence": stats.max_confidence,
                        "min_confidence": stats.min_confidence,
                        "avg_confidence": stats.avg_confidence,
                        "total_objects": stats.total_objects
                    }
                
                # 获取车辆详情
                vehicle_details = []
                seen_vehicle_types = set()
                for detail in task.vehicle_details:
                    if detail.vehicle_type not in seen_vehicle_types:
                        vehicle_details.append({
                            "vehicle_type": detail.vehicle_type,
                            "count": detail.count
                        })
                        seen_vehicle_types.add(detail.vehicle_type)
                
                # 从input_file路径中提取文件名
                filename = os.path.basename(task.input_file)
                
                # 构建图片URL
                processed_image_url = None
                if task.output_file and os.path.exists(task.output_file):
                    # 将绝对路径转换为相对URL路径
                    relative_path = os.path.relpath(task.output_file, UPLOAD_DIR)
                    processed_image_url = f"/uploads/{relative_path}"
                    logger.info(f"Generated image URL: {processed_image_url}")
                    logger.info(f"Full file path: {task.output_file}")
                    logger.info(f"File exists: {os.path.exists(task.output_file)}")
                
                result = {
                    "task_id": task.id,
                    "filename": filename,
                    "media_type": task.media_type,
                    "status": task.status,
                    "created_at": task.created_at.isoformat(),
                    "pedestrian_count": task.pedestrian_count,
                    "vehicle_count": task.vehicle_count,
                    "processing_time": task.processing_time,
                    "detection_details": json.loads(detection_record.detection_results) if detection_record else [],
                    "vehicle_details": vehicle_details,
                    "statistics": statistics,
                    "result_url": f"/result/{task.id}",
                    "processed_image_url": processed_image_url,
                    "confidence_threshold": task.confidence_threshold,
                    "iou_threshold": task.iou_threshold,
                    "max_detections": task.max_detections
                }
                results.append(result)
                successful_tasks += 1
                
            except Exception as e:
                logger.error(f"Error processing task {task_id}: {str(e)}")
                errors.append(f"Error processing task {task_id}: {str(e)}")
                failed_tasks += 1
                continue
        
        return {
            "success": True,
            "total_tasks": total_tasks,
            "successful_tasks": successful_tasks,
            "failed_tasks": failed_tasks,
            "results": results,
            "errors": errors
        }
        
    except Exception as e:
        logger.error(f"Error in result endpoint: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# 历史记录分组查看
@app.get("/history/{user_id}")
async def get_history(user_id: int, db: Session = Depends(get_db)):
    """获取用户历史记录"""
    logger.info(f"获取用户 {user_id} 的历史记录")
    
    # 查询数据库
    histories = db.query(History).filter(History.user_id == user_id, History.is_deleted == False).all()
    
    # 格式化返回结果
    result = []
    for history in histories:
        # 检查输出文件路径
        output_file = history.output_file
        if output_file:
            if not os.path.exists(output_file):
                logger.warning(f"警告: 文件不存在 {output_file}")
                output_file = None
            else:
                # 确保路径格式正确
                output_file = normalize_path(output_file)
        
        history_info = {
            "id": history.id,
            "user_id": history.user_id,
            "input_file": history.input_file,
            "output_file": output_file,
            "pedestrian_count": history.pedestrian_count,
            "vehicle_count": history.vehicle_count,
            "confidence_threshold": history.confidence_threshold,
            "iou_threshold": history.iou_threshold,
            "max_detections": history.max_detections,
            "media_type": history.media_type,
            "created_at": history.created_at.isoformat()
        }
        
        # 获取车辆详情
        vehicle_details = []
        for detail in history.vehicle_details:
            vehicle_details.append({
                "vehicle_type": detail.vehicle_type,
                "count": detail.count
            })
        history_info["vehicle_details"] = vehicle_details
        
        result.append(history_info)
    
    logger.info(f"返回数据: {result}")
    return result

# 按时间分组查看历史记录
@app.get("/history/{user_id}/time/{time_period}")
def get_history_by_time(user_id: int, time_period: str, db: Session = Depends(get_db)):
    # 根据时间周期过滤历史记录
    if time_period == "today":
        histories = db.query(History).filter(History.user_id == user_id, History.created_at >= datetime.now().date()).all()
    elif time_period == "yesterday":
        histories = db.query(History).filter(History.user_id == user_id, History.created_at >= datetime.now().date() - timedelta(days=1)).all()
    elif time_period == "last_week":
        histories = db.query(History).filter(History.user_id == user_id, History.created_at >= datetime.now().date() - timedelta(days=7)).all()
    elif time_period == "this_month":
        histories = db.query(History).filter(History.user_id == user_id, History.created_at >= datetime.now().date() - timedelta(days=30)).all()
    else:
        raise HTTPException(status_code=400, detail="无效的时间周期")
    
    # 格式化返回结果
    result = []
    for history in histories:
        history_info = {
            "id": history.id,
            "user_id": history.user_id,
            "input_file": history.input_file,
            "output_file": history.output_file,
            "pedestrian_count": history.pedestrian_count,
            "vehicle_count": history.vehicle_count,
            "confidence_threshold": history.confidence_threshold,
            "iou_threshold": history.iou_threshold,
            "max_detections": history.max_detections,
            "media_type": history.media_type,
            "created_at": history.created_at.isoformat()
        }
        
        # 添加处理后的图片URL，使用标准化的路径
        if history.output_file and os.path.exists(history.output_file):
            history_info["processed_image_url"] = normalize_path(history.output_file)
        else:
            history_info["processed_image_url"] = None
        
        # 获取车辆详情
        vehicle_details = []
        for detail in history.vehicle_details:
            vehicle_details.append({
                "vehicle_type": detail.vehicle_type,
                "count": detail.count
            })
        history_info["vehicle_details"] = vehicle_details
        
        result.append(history_info)
    
    return result

# 按检测对象分组查看历史记录
@app.get("/history/{user_id}/object/{object_type}")
def get_history_by_object(user_id: int, object_type: str, db: Session = Depends(get_db)):
    # 根据检测对象过滤历史记录
    if object_type.lower() == "person":
        histories = db.query(History).filter(
            History.user_id == user_id,
            History.pedestrian_count > 0
        ).all()
    elif object_type.lower() == "vehicle":
        histories = db.query(History).filter(
            History.user_id == user_id,
            History.vehicle_count > 0
        ).all()
    else:
        raise HTTPException(status_code=400, detail="无效的对象类型，只支持 'person' 或 'vehicle'")
    
    # 格式化返回结果
    result = []
    for history in histories:
        history_info = {
            "id": history.id,
            "user_id": history.user_id,
            "input_file": history.input_file,
            "output_file": history.output_file,
            "pedestrian_count": history.pedestrian_count,
            "vehicle_count": history.vehicle_count,
            "confidence_threshold": history.confidence_threshold,
            "iou_threshold": history.iou_threshold,
            "max_detections": history.max_detections,
            "media_type": history.media_type,
            "created_at": history.created_at.isoformat()
        }
        
        # 添加处理后的图片URL，使用标准化的路径
        if history.output_file and os.path.exists(history.output_file):
            history_info["processed_image_url"] = normalize_path(history.output_file)
        else:
            history_info["processed_image_url"] = None
        
        # 获取车辆详情
        vehicle_details = []
        for detail in history.vehicle_details:
            vehicle_details.append({
                "vehicle_type": detail.vehicle_type,
                "count": detail.count
            })
        history_info["vehicle_details"] = vehicle_details
        
        result.append(history_info)
    
    return {
        "total": len(result),
        "histories": result
    }

# 按任务状态分组查看历史记录
@app.get("/history/{user_id}/status/{status}")
def get_history_by_status(user_id: int, status: str, db: Session = Depends(get_db)):
    # 根据任务状态过滤历史记录
    if status == "deleted":
        histories = db.query(History).filter(History.user_id == user_id, History.is_deleted == True).all()
    else:
        histories = db.query(History).filter(History.user_id == user_id, History.is_deleted == False).all()
    return histories

# 回收记录（移入回收站）
@app.put("/history/{history_id}/remove")
def remove_history(history_id: int, db: Session = Depends(get_db)):
    history = db.get(History, history_id)
    if not history:
        raise HTTPException(status_code=404, detail="记录不存在")
    history.is_deleted = True
    db.commit()
    return {"success": True}

# 查看回收站
@app.get("/recycle/{user_id}")
def get_recycle(user_id: int, db: Session = Depends(get_db)):
    histories = db.query(History).filter(History.user_id == user_id, History.is_deleted == True).all()
    return histories

# 恢复记录
@app.put("/recycle/{history_id}/restore")
def restore_history(history_id: int, db: Session = Depends(get_db)):
    history = db.get(History, history_id)
    if not history:
        raise HTTPException(status_code=404, detail="记录不存在")
    history.is_deleted = False
    db.commit()
    return {"success": True}

# 永久删除记录
@app.delete("/recycle/{history_id}")
def delete_history(history_id: int, db: Session = Depends(get_db)):
    history = db.get(History, history_id)
    if not history:
        raise HTTPException(status_code=404, detail="记录不存在")
    db.delete(history)
    db.commit()
    return {"success": True}

# 批量删除历史记录（移入回收站）
@app.delete("/history/batch-delete")
def batch_delete_history(history_ids: List[int], db: Session = Depends(get_db)):
    # 删除此接口，使用 /history/batch-operation 替代
    pass

# 批量移入回收站
@app.put("/history/batch-remove")
def batch_remove_history(history_ids: List[int], db: Session = Depends(get_db)):
    # 删除此接口，使用 /history/batch-operation 替代
    pass

# 按媒体类型分组查看历史记录
@app.get("/history/{user_id}/media/{media_type}")
def get_history_by_media_type(user_id: int, media_type: str, db: Session = Depends(get_db)):
    # 根据媒体类型过滤历史记录
    if media_type.lower() not in ["image", "video"]:
        raise HTTPException(status_code=400, detail="无效的媒体类型，只支持 'image' 或 'video'")
    
    histories = db.query(History).filter(
        History.user_id == user_id,
        History.media_type == media_type.lower(),
        History.is_deleted == False
    ).all()
    
    # 格式化返回结果
    result = []
    for history in histories:
        history_info = {
            "id": history.id,
            "user_id": history.user_id,
            "input_file": history.input_file,
            "output_file": history.output_file,
            "pedestrian_count": history.pedestrian_count,
            "vehicle_count": history.vehicle_count,
            "confidence_threshold": history.confidence_threshold,
            "iou_threshold": history.iou_threshold,
            "max_detections": history.max_detections,
            "media_type": history.media_type,
            "created_at": history.created_at.isoformat()
        }
        
        # 添加处理后的图片URL，使用标准化的路径
        if history.output_file and os.path.exists(history.output_file):
            history_info["processed_image_url"] = normalize_path(history.output_file)
        else:
            history_info["processed_image_url"] = None
        
        # 获取车辆详情
        vehicle_details = []
        for detail in history.vehicle_details:
            vehicle_details.append({
                "vehicle_type": detail.vehicle_type,
                "count": detail.count
            })
        history_info["vehicle_details"] = vehicle_details
        
        result.append(history_info)
    
    return {
        "total": len(result),
        "histories": result
    }

# 按媒体类型查看回收站记录
@app.get("/recycle/{user_id}/media/{media_type}")
def get_recycle_by_media_type(user_id: int, media_type: str, db: Session = Depends(get_db)):
    # 根据媒体类型过滤回收站记录
    if media_type.lower() not in ["image", "video"]:
        raise HTTPException(status_code=400, detail="无效的媒体类型，只支持 'image' 或 'video'")
    
    histories = db.query(History).filter(
        History.user_id == user_id,
        History.media_type == media_type.lower(),
        History.is_deleted == True
    ).all()
    
    # 格式化返回结果
    result = []
    for history in histories:
        history_info = {
            "id": history.id,
            "user_id": history.user_id,
            "input_file": history.input_file,
            "output_file": history.output_file,
            "pedestrian_count": history.pedestrian_count,
            "vehicle_count": history.vehicle_count,
            "confidence_threshold": history.confidence_threshold,
            "iou_threshold": history.iou_threshold,
            "max_detections": history.max_detections,
            "media_type": history.media_type,
            "created_at": history.created_at.isoformat()
        }
        
        # 获取车辆详情
        vehicle_details = []
        for detail in history.vehicle_details:
            vehicle_details.append({
                "vehicle_type": detail.vehicle_type,
                "count": detail.count
            })
        history_info["vehicle_details"] = vehicle_details
        
        result.append(history_info)
    
    return {
        "total": len(result),
        "histories": result
    }

# 统计分析API
@app.get("/statistics/{user_id}")
def get_statistics(user_id: int, db: Session = Depends(get_db)):
    # 获取用户的所有任务
    tasks = db.query(Task).filter(Task.user_id == user_id).all()
    
    # 计算总体统计
    total_tasks = len(tasks)
    completed_tasks = len([t for t in tasks if t.status == TASK_STATUS_COMPLETED])
    failed_tasks = len([t for t in tasks if t.status == TASK_STATUS_FAILED])
    
    # 计算媒体类型统计
    image_tasks = len([t for t in tasks if t.input_file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.gif'))])
    video_tasks = len([t for t in tasks if t.input_file.lower().endswith(('.mp4', '.avi', '.mov', '.mkv', '.wmv'))])
    
    # 计算检测对象统计
    total_pedestrians = sum(t.pedestrian_count for t in tasks)
    total_vehicles = sum(t.vehicle_count for t in tasks)
    
    # 获取最近7天的每日统计
    daily_stats = []
    for i in range(7):
        date = datetime.now().date() - timedelta(days=i)
        day_tasks = [t for t in tasks if t.created_at.date() == date]
        daily_stats.append({
            "date": date.strftime("%Y-%m-%d"),
            "task_count": len(day_tasks),
            "image_count": len([t for t in day_tasks if t.input_file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.gif'))]),
            "video_count": len([t for t in day_tasks if t.input_file.lower().endswith(('.mp4', '.avi', '.mov', '.mkv', '.wmv'))]),
            "pedestrian_count": sum(t.pedestrian_count for t in day_tasks),
            "vehicle_count": sum(t.vehicle_count for t in day_tasks)
        })
    
    # 获取车辆类型分布
    vehicle_types = {
        "car": 0,          # 小汽车
        "bus": 0,          # 公交车
        "truck": 0,        # 卡车
        "motorcycle": 0,   # 摩托车
        "bicycle": 0,      # 自行车
        "van": 0,          # 面包车
        "trailer": 0,      # 拖车
        "ambulance": 0,    # 救护车
        "fire_truck": 0,   # 消防车
        "police_car": 0    # 警车
    }
    
    for task in tasks:
        for detail in task.vehicle_details:
            if detail.vehicle_type in vehicle_types:
                vehicle_types[detail.vehicle_type] += detail.count
    
    return {
        "overall": {
            "total_tasks": total_tasks,
            "completed_tasks": completed_tasks,
            "failed_tasks": failed_tasks,
            "image_tasks": image_tasks,
            "video_tasks": video_tasks,
            "total_pedestrians": total_pedestrians,
            "total_vehicles": total_vehicles
        },
        "daily_stats": daily_stats,
        "vehicle_distribution": vehicle_types,
        "media_type_distribution": {
            "images": image_tasks,
            "videos": video_tasks,
            "image_percentage": round(image_tasks / total_tasks * 100 if total_tasks > 0 else 0, 2),
            "video_percentage": round(video_tasks / total_tasks * 100 if total_tasks > 0 else 0, 2)
        }
    }

# 获取处理结果文件
@app.get("/result/{task_id}")
async def get_result_file(task_id: str, db: Session = Depends(get_db)):
    # 获取任务信息
    task = db.query(Task).filter(Task.id == task_id).first()
    if not task:
        raise HTTPException(status_code=404, detail="任务不存在")
    
    # 检查任务状态
    if task.status != TASK_STATUS_COMPLETED:
        raise HTTPException(status_code=400, detail="任务尚未完成处理")
    
    # 获取输出文件路径
    output_path = task.output_file
    if not output_path or not os.path.exists(output_path):
        raise HTTPException(status_code=404, detail="处理结果文件不存在")
    
    # 返回文件
    return FileResponse(
        path=output_path,
        filename=os.path.basename(output_path),
        media_type="image/jpeg" if task.media_type == "image" else "video/mp4"
    )

# 添加静态文件服务
app.mount("/uploads", StaticFiles(directory=UPLOAD_DIR), name="uploads")

def find_free_port(start_port=8000, max_port=8999):
    """查找可用的端口"""
    for port in range(start_port, max_port + 1):
        with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
            try:
                sock.bind(('0.0.0.0', port))
                return port
            except OSError:
                continue
    raise RuntimeError(f"Could not find a free port between {start_port} and {max_port}")

@app.get("/api/tasks/{task_id}/result")
async def get_task_result(task_id: str):
    """获取任务结果"""
    try:
        # 使用本地定义的 get_task_result 函数，而不是从 tasks 模块导入
        result = get_task_result(task_id)
        return result
    except HTTPException as e:
        raise e
    except Exception as e:
        logger.error(f"获取任务结果失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# 将 get_task_result 函数移到这个位置
def get_task_result(task_id: str) -> Dict:
    """获取任务结果"""
    db = SessionLocal()
    try:
        task = db.query(Task).filter(Task.id == task_id).first()
        if not task:
            raise HTTPException(status_code=404, detail="Task not found")
        
        result = {
            "task_id": task.id,
            "status": task.status,
            "output_file": task.output_file,
            "error": task.error,
            "confidence_threshold": task.confidence_threshold,
            "iou_threshold": task.iou_threshold,
            "max_detections": task.max_detections,
            "pedestrian_count": task.pedestrian_count,
            "vehicle_count": task.vehicle_count,
            "processing_time": task.processing_time
        }
        return result
    finally:
        db.close()

if __name__ == "__main__":
    import uvicorn
    try:
        # 尝试使用默认端口
        port = 8000
        uvicorn.run(app, host="0.0.0.0", port=port)
    except OSError as e:
        if "address already in use" in str(e).lower():
            # 如果默认端口被占用，查找可用端口
            port = find_free_port()
            logger.info(f"Port 8000 is in use, switching to port {port}")
            uvicorn.run(app, host="0.0.0.0", port=port)
        else:
            raise
